import json
import os
import pymysql
from dbutils.pooled_db import PooledDB
import threading
from typing import Dict, Optional

class MCPConfig:
    def __init__(self):
        self.config_path = os.path.join(os.path.dirname(__file__), 'config', 'mcp.json')
        self.config = self._load_config()

    def _load_config(self):
        """加载 MCP 配置文件"""
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"加载 MCP 配置文件失败: {str(e)}")
            return None

    def get_server_config(self, server_key: str = None):
        """获取指定服务器的配置"""
        if not self.config:
            return None

        if not server_key:
            server_key = self.config.get('defaultServer', 'default')

        return self.config.get('servers', {}).get(server_key)

    def get_global_config(self):
        """获取全局配置"""
        return self.config.get('globalConfig', {}) if self.config else {}

class MySQLPoolManager:
    _instance = None
    _lock = threading.Lock()
    _pools: Dict[str, PooledDB] = {}

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._mcp_config = MCPConfig()
        return cls._instance

    def get_pool(self, pool_key: str = None) -> Optional[PooledDB]:
        """获取指定的连接池"""
        if pool_key is None:
            pool_key = self._mcp_config.config.get('defaultServer', 'default')
        return self._pools.get(pool_key)

    def create_pool(self, pool_key: str = None, **kwargs) -> PooledDB:
        """创建新的连接池"""
        if pool_key is None:
            pool_key = self._mcp_config.config.get('defaultServer', 'default')

        if pool_key in self._pools:
            return self._pools[pool_key]

        with self._lock:
            if pool_key not in self._pools:
                # 获取服务器配置
                server_config = self._mcp_config.get_server_config(pool_key)
                if not server_config:
                    raise Exception(f"未找到服务器配置: {pool_key}")

                # 获取全局配置
                global_config = self._mcp_config.get_global_config()

                # 合并配置
                pool_config = {
                    'creator': pymysql,
                    'host': server_config.get('host'),
                    'port': server_config.get('port'),
                    'user': server_config.get('user'),
                    'password': server_config.get('password'),
                    'database': server_config.get('database'),
                    'charset': global_config.get('charset', 'utf8mb4'),
                    'maxconnections': server_config.get('poolConfig', {}).get('maxConnections', 10),
                    'mincached': server_config.get('poolConfig', {}).get('minConnections', 2),
                    'maxcached': server_config.get('poolConfig', {}).get('maxConnections', 10),
                    'maxshared': server_config.get('poolConfig', {}).get('maxConnections', 10),
                    'blocking': True,
                    'maxusage': None,
                    'setsession': [],
                    'ping': 1,
                    'cursorclass': pymysql.cursors.DictCursor
                }

                # 更新自定义配置
                pool_config.update(kwargs)

                # 创建连接池
                self._pools[pool_key] = PooledDB(**pool_config)

        return self._pools[pool_key]

    def close_pool(self, pool_key: str = None):
        """关闭指定的连接池"""
        if pool_key is None:
            pool_key = self._mcp_config.config.get('defaultServer', 'default')

        if pool_key in self._pools:
            with self._lock:
                if pool_key in self._pools:
                    del self._pools[pool_key]

    def close_all_pools(self):
        """关闭所有连接池"""
        with self._lock:
            self._pools.clear()

class DatabaseConnection:
    def __init__(self, pool_key: str = None):
        self.pool_key = pool_key
        self.pool_manager = MySQLPoolManager()
        self.connection = None
        self.cursor = None

    def __enter__(self):
        pool = self.pool_manager.get_pool(self.pool_key)
        if pool is None:
            # 如果连接池不存在，尝试创建
            pool = self.pool_manager.create_pool(self.pool_key)
        self.connection = pool.connection()
        self.cursor = self.connection.cursor()
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()

# 初始化数据库连接池
def initialize_db_pools():
    """初始化所有配置的数据库连接池"""
    pool_manager = MySQLPoolManager()
    mcp_config = MCPConfig()

    # 获取所有服务器配置
    servers = mcp_config.config.get('servers', {})

    # 为每个服务器创建连接池
    for server_key in servers:
        pool_manager.create_pool(server_key)
        print(f"已初始化连接池: {server_key}")

# 便捷的数据库操作方法
def execute_query(sql: str, params: tuple = None, pool_key: str = None):
    """执行查询"""
    with DatabaseConnection(pool_key) as cursor:
        cursor.execute(sql, params)
        return cursor.fetchall()

def execute_update(sql: str, params: tuple = None, pool_key: str = None):
    """执行更新"""
    with DatabaseConnection(pool_key) as cursor:
        try:
            result = cursor.execute(sql, params)
            cursor.connection.commit()
            return result
        except Exception as e:
            cursor.connection.rollback()
            raise e
